package com.shiyi.shiyiaiagent.advisor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.MessageAggregator;
import reactor.core.publisher.Flux;

@Slf4j
public class MyLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

  private AdvisedRequest before(AdvisedRequest request){
   log.info("AI request:{}",request.userText());
   return request;
  }

  private void observeAfter(AdvisedResponse response){
    log.info("AI response:{}",response.response().getResult().getOutput().getText());
  }


  @Override
  public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
    advisedRequest = this.before(advisedRequest);
    AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);
    this.observeAfter(advisedResponse);
    return advisedResponse;
  }

  @Override
  public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {
    advisedRequest = this.before(advisedRequest);
    Flux<AdvisedResponse> advisedResponseFlux = chain.nextAroundStream(advisedRequest);
    return new MessageAggregator().aggregateAdvisedResponse(advisedResponseFlux,this::observeAfter);
  }

  @Override
  public String getName() {
    return this.getClass().getSimpleName();
  }

  @Override
  public int getOrder() {
    return 0;
  }
}
